﻿#pragma once

#include <learn/math/math.h>
#include <learn/math/matrix.h>

namespace xi
{

namespace math
{

/**
 * @brief 获得导函数
 *
 * @tparam T
 * @param[in] f
 * @return std::function<T(double)>
 */
template <class T> inline std::function<T(double)> differential(std::function<T(double)> f)
{
    static constexpr double h = Numeric::distance;
    return [=](double x) { return (f(x + h) - f(x - h)) / (2 * h); };
}

/**
 * @brief 多元函数的 Jacobian 矩阵
 *
 * @tparam T
 * @param[in] f
 * @param[in] x
 * @return MatrixXd
 */
template <class T> inline MatrixXd jacobian(std::function<T(VectorXd)> f, VectorXd x)
{
    // 巨坑：有些版本编译器不支持将 static_assert(false) 写在 else 里
    static constexpr double h = Numeric::distance;
    static_assert(std::is_same_v<T, VectorXd> || std::is_same_v<T, double>, "Unsupported type.");

    T y = f(x);
    std::size_t m = 0;
    if constexpr (std::is_same_v<T, VectorXd>)
        m = y.size();
    else if constexpr (std::is_same_v<T, double>)
        m = 1;

    std::size_t n = x.size();
    MatrixXd d(m, n);
    for (std::size_t i = 0; i < n; i++)
    {
        VectorXd y1 = x, y2 = x;
        y1[i] += h, y2[i] -= h;

        // 中心差分向量
        T dp = (f(y1) - f(y2)) / (2 * h);
        if constexpr (std::is_same_v<T, VectorXd>)
        {
            for (std::size_t j = 0; j < m; j++)
                d(j, i) = dp[j];
        }
        else if constexpr (std::is_same_v<T, double>)
            d(0, i) = dp;
    }
    return d;
}

/**
 * @brief 多元函数的梯度向量
 *
 * @tparam T
 * @param[in] f
 * @param[in] x
 * @return MatrixXd
 */
template <class T> inline MatrixXd gradient(std::function<T(VectorXd)> f, VectorXd x)
{
    return jacobian(f, x).transpose();
}

/**
 * @brief 多元函数的 Hessian 矩阵
 *
 * @param[in] f
 * @param[in] x
 * @return MatrixXd
 */
inline MatrixXd hessian(std::function<double(VectorXd)> f, VectorXd x)
{
    static constexpr double h = Numeric::distance;

    Matrix d = jacobian(f, x);
    std::size_t n = x.size();
    Matrix H(n, n);
    for (std::size_t i = 0; i < n; i++)
    {
        VectorXd y1 = x, y2 = x;
        y1[i] += h, y2[i] -= h;

        Matrix dp = (jacobian(f, y1) - jacobian(f, y2)) / (2 * h);
        for (std::size_t j = 0; j < n; j++)
            H(i, j) = dp(0, j);
    }
    return H;
}

} // namespace math

} // namespace xi
